/*
   Copyright 2009 Abdulla Abdurakhmanov (abdulla@latestbit.com)
   Original sources are available at www.latestbit.com

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package org.latestbit.smartobjects.network;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.security.GeneralSecurityException;
import java.security.PublicKey;
import java.util.Date;

import javax.crypto.SecretKey;

import org.latestbit.smartobjects.coders.*;
import org.latestbit.smartobjects.encrypters.IAsymmetricDataEncrypter;
import org.latestbit.smartobjects.encrypters.ISymmetricDataEncrypter;
import org.latestbit.smartobjects.exceptions.*;
import org.latestbit.smartobjects.io.ResizableByteBuffer;

public class StdHandshakeManager implements IHandshakeManager {

	private static final String SO_STD_HANDSHAKE_TYPE = "SOSH";

	protected enum HandshakeStatus {
		WAITING_TO_BEGIN_HANDSHAKE,
		WAITING_TO_PUBLIC_KEY,
		WAITING_TO_AUTH,
		WAITING_TO_SYMMETRIC_KEY,
		HANDSHAKED,
		ERROR
	}
	
	protected HandshakeStatus handshakeStatus = HandshakeStatus.WAITING_TO_BEGIN_HANDSHAKE;
	private ResizableByteBuffer handshakeBuffer = new ResizableByteBuffer(false, 128, 1024, 128);
	private ResizableByteBuffer bufferToCrypt = new ResizableByteBuffer(false, 128, 1024, 128);

	private IHandshakeManagerFactory factory;
	private IConnection connection;
	private PublicKey remotePublicKey;
	private IAsymmetricDataEncrypter remoteAssymetricEncrypter;

	public StdHandshakeManager( IHandshakeManagerFactory factory, IConnection connection ) {		
		this.factory = factory;
		this.connection = connection;
	}
		
	@Override
	public void reset( ) {
		this.handshakeStatus = HandshakeStatus.WAITING_TO_BEGIN_HANDSHAKE;
		this.remotePublicKey = null;
	}	

	@Override
	public boolean isHandshaked() {
		return handshakeStatus == HandshakeStatus.HANDSHAKED;
	}

	@Override
	public void beginHandshake( ) throws HandshakeException {
		handshakeStatus = HandshakeStatus.WAITING_TO_PUBLIC_KEY;
		handshakeBuffer.getBuffer().clear();
		HandshakeRequest request = new HandshakeRequest(handshakeBuffer, this.factory.getBus().createDataTypesCoder());
		sendHandshakeBegin(request);
		// sendPublicKey(request );
		if(this.factory.getHandshakeManagerListener()!=null) {
			this.factory.getHandshakeManagerListener().onHandshakeBegin( this );
		}
	}
	
	protected void sendHandshakeBegin(StdNetworkCommand command ) throws HandshakeException {
		IDataTypesCoder coder = command.getDataTypesCoder();
		try {
			command.beginSerialize();
			coder.encodeString(SO_STD_HANDSHAKE_TYPE, command.getBuffer());
			command.finishSerialize();
			connection.send( command.getBuffer().getBuffer() );
		} 
		catch (Exception ex) {
			throw new HandshakeException (ex);
		}
	}
	

	protected void sendPublicKey(StdNetworkCommand command ) throws HandshakeException {
		IAsymmetricDataEncrypter asymmetricEncrypter = factory.getAsymmetricDataEncrypter();
		try {
			command.beginSerialize();
			asymmetricEncrypter.encodePublicKey( this.factory.getKeyPair().getPublic(),  command.getBuffer() );
			command.finishSerialize();
			connection.send( command.getBuffer().getBuffer() );
		} 
		catch (Exception ex) {
			throw new HandshakeException (ex);
		}
	}
	
	protected void setupRemoteEncrypter(ByteBuffer buffer) throws GeneralSecurityException {
		this.remotePublicKey = factory.getAsymmetricDataEncrypter().decodePublicKey ( buffer );
		this.remoteAssymetricEncrypter = this.factory.getBus().getDataEncrypterFactory().createAsymmetricEncrypter();
		this.remoteAssymetricEncrypter.initCipher( this.remotePublicKey );
	}
	
	@Override
	public void onHandshakeReceived( ByteBuffer buffer) throws HandshakeException {
		
		IDataTypesCoder coder = this.factory.getBus().createDataTypesCoder();
		StdNetworkCommand command = new StdNetworkCommand ( buffer, coder );
		
		try {
			command.deserialize();
			if(command.getCommandId() == HandshakeResponse.COMMAND_ID) {
				if( !this.connection.isAcceptorConnection()) {
					HandshakeResponse response = new HandshakeResponse ( command );
					response.deserialize();
					onHandshakeResponseReceived ( response );
				}
				else
					throw new HandshakeException("Unexpected handshake response is received on acceptor connection");
			}
			else
				if(command.getCommandId() == HandshakeRequest.COMMAND_ID) {
					if( this.connection.isAcceptorConnection()) {
						HandshakeRequest request = new HandshakeRequest ( command );
						request.deserialize();
						try {
							onHandshakeRequestReceived ( request );
						}
						catch(Exception ex) {
							sendErrorHandshakeResponse ( request, ex );
						}
					}
					else
						throw new HandshakeException("Unexpected handshake request is received on connector");						
				}
		} 
		catch (Exception ex) {
			throw new HandshakeException (ex);
		}
	}

	protected void onHandshakeRequestReceived ( HandshakeRequest request) throws HandshakeException, GeneralSecurityException, InvalidCredentialsException, SmartBusSystemException, IOException {
		if ( handshakeStatus == HandshakeStatus.WAITING_TO_BEGIN_HANDSHAKE ) {
			onReceiveClientBeginHandshake ( request );
		}
		else
		if ( handshakeStatus == HandshakeStatus.WAITING_TO_SYMMETRIC_KEY ) {
			onReceiveClientSymmetricKey ( request );
		}
		else
		if ( handshakeStatus == HandshakeStatus.WAITING_TO_AUTH ) {
			onReceiveClientAuth ( request );
		}
	}

	protected void onReceiveClientBeginHandshake( HandshakeRequest request ) throws HandshakeException, GeneralSecurityException {		
		handshakeBuffer.getBuffer().clear();
		handshakeBuffer.put( request.getBuffer().getBuffer() );
		handshakeBuffer.getBuffer().flip();
		
		IDataTypesCoder coder = request.getDataTypesCoder();
		try {
			String handshakeType = coder.decodeString ( handshakeBuffer.getBuffer() );
			if(!handshakeType.equalsIgnoreCase(SO_STD_HANDSHAKE_TYPE))
				throw new HandshakeException("Unknown handshake type received: "+handshakeType);
		}
		catch(DataTypesCoderException ex) {
			throw new HandshakeException(ex);
		}
		
		// Sending server public key			
		handshakeBuffer.getBuffer().clear();
		sendPublicKey( request.createResponse(handshakeBuffer) );
		handshakeStatus = HandshakeStatus.WAITING_TO_SYMMETRIC_KEY;
		
		if(this.factory.getHandshakeManagerListener()!=null) {
			this.factory.getHandshakeManagerListener().onHandshakeBegin( this );
		}		
	}
	
	protected void onReceiveClientSymmetricKey( HandshakeRequest request ) throws SmartBusSystemException, GeneralSecurityException, IOException {
		// Decrypt symmetric key
		bufferToCrypt.getBuffer().clear();
		handshakeBuffer.getBuffer().clear();
		handshakeBuffer.put( request.getBuffer().getBuffer() );
		handshakeBuffer.getBuffer().flip();
		int encryptDataLen = request.getCoder().decodeEncryptedDataTypeBegin( handshakeBuffer.getBuffer() );
		int saveLimit = handshakeBuffer.getBuffer().limit();
		handshakeBuffer.getBuffer().limit(handshakeBuffer.getBuffer().position() + encryptDataLen);		
		
		this.factory.getAsymmetricDataEncrypter().decrypt( handshakeBuffer.getBuffer(), bufferToCrypt );
		bufferToCrypt.getBuffer().flip();
		handshakeBuffer.getBuffer().limit(saveLimit);
		request.getCoder().decodeEncryptedDataTypeEnd( handshakeBuffer.getBuffer() );
	
		ISymmetricDataEncrypter remoteSymmetricEncrypter = this.factory.getBus().getDataEncrypterFactory().createSymmetricEncrypter();
		remoteSymmetricEncrypter.initCipher( remoteSymmetricEncrypter.decodeSecretKey(bufferToCrypt.getBuffer()) );		
		
		handshakeBuffer.getBuffer().clear();
		bufferToCrypt.getBuffer().clear();
		
		// Creating server symmetric encrypter & key
		ISymmetricDataEncrypter connectionEncrypter = this.factory.getBus().getDataEncrypterFactory().createSymmetricEncrypter();
		SecretKey sessionSymmetricKey = connectionEncrypter.generateSecret();
		connectionEncrypter.initCipher( sessionSymmetricKey );		
		connectionEncrypter.encodeSecretKey( sessionSymmetricKey, bufferToCrypt );
		connection.setDataEncrypter(connectionEncrypter);
		
		// Sending server session info & session symmetric key encoded by remote symmetric key
		HandshakeResponse response = request.createResponse(handshakeBuffer);		
		response.beginSerialize();
		
		int encBegPos = response.getCoder().encodeEncryptedDataTypeBegin( handshakeBuffer );
		bufferToCrypt.getBuffer().flip();
		remoteSymmetricEncrypter.encrypt( bufferToCrypt.getBuffer(), handshakeBuffer );		
		response.getCoder().encodeEncryptedDataTypeEnd (encBegPos, handshakeBuffer );		
		
		response.finishSerialize();
		connection.send( response.getBuffer().getBuffer() );
		
		
		handshakeStatus = HandshakeStatus.WAITING_TO_AUTH;
	}

	protected void onReceiveClientAuth( HandshakeRequest request ) throws SmartBusSystemException, GeneralSecurityException, IOException {

		// Decrypt credentials
		bufferToCrypt.getBuffer().clear();
		handshakeBuffer.getBuffer().clear();
		handshakeBuffer.put( request.getBuffer().getBuffer() );
		handshakeBuffer.getBuffer().flip();
		int encryptDataLen = request.getCoder().decodeEncryptedDataTypeBegin( handshakeBuffer.getBuffer() );
		int saveLimit = handshakeBuffer.getBuffer().limit();
		handshakeBuffer.getBuffer().limit(handshakeBuffer.getBuffer().position() + encryptDataLen);		
		
		connection.getDataEncrypter().decrypt( handshakeBuffer.getBuffer(), bufferToCrypt );
		bufferToCrypt.getBuffer().flip();
		handshakeBuffer.getBuffer().limit(saveLimit);
		request.getCoder().decodeEncryptedDataTypeEnd( handshakeBuffer.getBuffer() );

		request.getCoder().decodeDateTime( bufferToCrypt.getBuffer() );
		// Decoding crentials & auth
		factory.getConnectionAuthenticator().decodeCredentials( this, request.getCoder(), bufferToCrypt.getBuffer() );
				
		handshakeBuffer.getBuffer().clear();
		bufferToCrypt.getBuffer().clear();
		

		// Sending server OK for auth
		HandshakeResponse response = request.createResponse(handshakeBuffer);		
		response.beginSerialize();		
		response.finishSerialize();
		connection.send( response.getBuffer().getBuffer() );
		
		handshakeStatus = HandshakeStatus.HANDSHAKED;
		if(this.factory.getHandshakeManagerListener()!=null) {
			this.factory.getHandshakeManagerListener().onHandshakeComplete( this );
		}
		
	}

	protected void sendErrorHandshakeResponse(HandshakeRequest request, Exception ex) throws SmartBusSystemException, IOException {
		handshakeBuffer.getBuffer().clear();
		bufferToCrypt.getBuffer().clear();
	
		HandshakeResponse response = request.createResponse(handshakeBuffer);
		response.setError(true);
		response.setErrorMessage(ex.toString());
		response.beginSerialize();
		response.finishSerialize();
		connection.send( response.getBuffer().getBuffer() );
		// Sleeping while a packet arrived before this connection will closed. 
		try {
			Thread.sleep(1000);
		}
		catch (InterruptedException e) {};
		handshakeStatus = HandshakeStatus.ERROR;
		
		if(this.factory.getHandshakeManagerListener()!=null) {
			this.factory.getHandshakeManagerListener().onHandshakeError( this, ex );
		}
		
	}

	protected void onHandshakeResponseReceived ( HandshakeResponse response ) throws GeneralSecurityException, SmartBusSystemException, IOException {
		if(!response.isError()) {
			if ( handshakeStatus == HandshakeStatus.WAITING_TO_PUBLIC_KEY ) {
				onReceiveServerPublicKey (  response );
			}
			else
			if ( handshakeStatus == HandshakeStatus.WAITING_TO_SYMMETRIC_KEY ) {
				onReceiveServerSymmetricKey (  response );
			}			
			else
			if ( handshakeStatus == HandshakeStatus.WAITING_TO_AUTH ) {
				onReceiveServerFinishHandshake (  response );
			}
		}
		else {
			onReceiveHandshakeError(response);
		}
	}
	

	protected void onReceiveServerPublicKey( HandshakeResponse response) throws GeneralSecurityException, SmartBusSystemException, IOException {		
		// Decode server public key & setup encrypter
		setupRemoteEncrypter( response.getBuffer().getBuffer() );
		bufferToCrypt.getBuffer().clear();		

		ISymmetricDataEncrypter temporaryConnectionEncrypter = this.factory.getBus().getDataEncrypterFactory().createSymmetricEncrypter();
		SecretKey sessionSymmetricKey = temporaryConnectionEncrypter.generateSecret();
		temporaryConnectionEncrypter.initCipher( sessionSymmetricKey );		
		temporaryConnectionEncrypter.encodeSecretKey( sessionSymmetricKey, bufferToCrypt );
		connection.setDataEncrypter(temporaryConnectionEncrypter);
		
		handshakeBuffer.getBuffer().clear();
		
		HandshakeRequest request = new HandshakeRequest( handshakeBuffer, response.getDataTypesCoder());
		request.beginSerialize();
		
		int encBegPos = request.getCoder().encodeEncryptedDataTypeBegin( handshakeBuffer );
		bufferToCrypt.getBuffer().flip();
		this.remoteAssymetricEncrypter.encrypt( bufferToCrypt.getBuffer(), handshakeBuffer );		
		request.getCoder().encodeEncryptedDataTypeEnd (encBegPos, handshakeBuffer );
		
		request.finishSerialize();
		connection.send( request.getBuffer().getBuffer() );
		
		this.handshakeStatus = HandshakeStatus.WAITING_TO_SYMMETRIC_KEY;
	}
	
	protected void onReceiveServerSymmetricKey( HandshakeResponse response) throws GeneralSecurityException, SmartBusSystemException, IOException {

		// Get the server symmetric key
		handshakeBuffer.getBuffer().clear();
		bufferToCrypt.getBuffer().clear();

		handshakeBuffer.put( response.getBuffer().getBuffer() );
		handshakeBuffer.getBuffer().flip();
		int encryptDataLen = response.getCoder().decodeEncryptedDataTypeBegin( handshakeBuffer.getBuffer() );
		int saveLimit = handshakeBuffer.getBuffer().limit();
		handshakeBuffer.getBuffer().limit(handshakeBuffer.getBuffer().position() + encryptDataLen);		
		
		connection.getDataEncrypter().decrypt( handshakeBuffer.getBuffer(), bufferToCrypt );
		bufferToCrypt.getBuffer().flip();
		handshakeBuffer.getBuffer().limit(saveLimit);
		response.getCoder().decodeEncryptedDataTypeEnd( handshakeBuffer.getBuffer() );
	
		ISymmetricDataEncrypter connectionSymmetricEncrypter = this.factory.getBus().getDataEncrypterFactory().createSymmetricEncrypter();
		connectionSymmetricEncrypter.initCipher( connectionSymmetricEncrypter.decodeSecretKey(bufferToCrypt.getBuffer()) );
		connection.setDataEncrypter(connectionSymmetricEncrypter);
		
		// Send credentials
		handshakeBuffer.getBuffer().clear();
		bufferToCrypt.getBuffer().clear();
		response.getCoder().encodeDateTime(new Date(), bufferToCrypt );
		factory.getConnectionAuthenticator().encodeCredentials( this, response.getCoder(), bufferToCrypt );		
		
		HandshakeRequest request = new HandshakeRequest( handshakeBuffer, response.getDataTypesCoder());
		request.beginSerialize();
		
		int encBegPos = request.getCoder().encodeEncryptedDataTypeBegin( handshakeBuffer );
		bufferToCrypt.getBuffer().flip();
		connection.getDataEncrypter().encrypt( bufferToCrypt.getBuffer(), handshakeBuffer );		
		request.getCoder().encodeEncryptedDataTypeEnd (encBegPos, handshakeBuffer );
		
		request.finishSerialize();
		connection.send( request.getBuffer().getBuffer() );
		
		this.handshakeStatus = HandshakeStatus.WAITING_TO_AUTH;
	}	

	protected void onReceiveServerFinishHandshake(HandshakeResponse response) throws DataTypesCoderException, GeneralSecurityException {
		
		handshakeBuffer.getBuffer().clear();
		bufferToCrypt.getBuffer().clear();
		handshakeStatus = HandshakeStatus.HANDSHAKED;

		if(this.factory.getHandshakeManagerListener()!=null) {
			this.factory.getHandshakeManagerListener().onHandshakeComplete( this );
		}
		
	}

	protected void onReceiveHandshakeError(HandshakeResponse response) {
		handshakeStatus = HandshakeStatus.ERROR;
		if(this.factory.getHandshakeManagerListener()!=null) {
			this.factory.getHandshakeManagerListener().onHandshakeError( this, new HandshakeException(response.getErrorMessage()) );
		}					
	}

	@Override
	public IConnection getConnection() {
		return this.connection;
	}

}
